/*++
THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
PARTICULAR PURPOSE.

Copyright (C) Microsoft Corporation.  All rights reserved.
*/

#include <windows.h>
#include <WinCrypt.h>
#include <tchar.h>
#include <stdio.h>

const int IN_BUFFER_SIZE    = 64 * 1024;
// OUT_BUFFER_SIZE is 8 bytes larger than IN_BUFFER_SIZE
// When CALG_RC2 algorithm is used, encrypted data
// will be 8 bytes larger than IN_BUFFER_SIZE
const int OUT_BUFFER_SIZE   = IN_BUFFER_SIZE + 8; // extra padding

void PrintUsage()
{
    _tprintf(_T("SessionKey <password> [</e>|</d>] <InputFile> <OutputFile>\n"));
    _tprintf(_T("/e for Encryption\n"));
    _tprintf(_T("/d for Decryption\n"));
}

void _tmain(int argc, TCHAR *argv[])
{
    BOOL fResult = FALSE;
    HCRYPTPROV hProv = NULL;
    HCRYPTHASH hHash = NULL;
    HCRYPTKEY hSessionKey = NULL;
    HANDLE hInFile = INVALID_HANDLE_VALUE;
    HANDLE hOutFile = INVALID_HANDLE_VALUE;
    BOOL fEncrypt = FALSE;
    BOOL finished = FALSE;
    BYTE pbBuffer[OUT_BUFFER_SIZE];
    DWORD dwByteCount = 0;
    DWORD dwBytesWritten = 0;

    if (argc != 5) {
        PrintUsage();
        return;
    }

    __try {
        /* Check whether the action to be performed is encrypt or decrypt */
        if (_tcsicmp(argv[2], _T("/e")) == 0) {
            fEncrypt = TRUE;
        } else if (_tcsicmp(argv[2], _T("/d")) == 0) {
            fEncrypt = FALSE;
        } else {
            PrintUsage();
            return;
        }

        // Open the input file to be encrypted or decrypted
        hInFile = CreateFile(argv[3],
                             GENERIC_READ,
                             0,
                             NULL,
                             OPEN_EXISTING,
                             FILE_ATTRIBUTE_NORMAL,
                             NULL);

        if (hInFile == INVALID_HANDLE_VALUE) {
            _tprintf(_T("CreateFile failed with %d\n"), GetLastError());
            __leave;
        }

        // Open the output file to write the encrypted or decrypted data
        hOutFile = CreateFile(argv[4],
                              GENERIC_WRITE,
                              0,
                              NULL,
                              CREATE_ALWAYS,
                              FILE_ATTRIBUTE_NORMAL,
                              NULL);

        if (hOutFile == INVALID_HANDLE_VALUE) {
            _tprintf(_T("CreateFile failed with %d\n"), GetLastError());
            __leave;
        }

        // Acquire a handle to MS_DEF_PROV using CRYPT_VERIFYCONTEXT for dwFlags
        // parameter as we are going to do only session key encryption or decryption
        fResult = CryptAcquireContext(&hProv,
                                      NULL,
                                      MS_DEF_PROV,
                                      PROV_RSA_FULL,
                                      CRYPT_VERIFYCONTEXT);

        if (!fResult) {
            _tprintf(_T("CryptAcquireContext failed with %X\n"), GetLastError());
            __leave;
        }

        fResult = CryptCreateHash(hProv, CALG_SHA1, 0, 0, &hHash);

        if (!fResult) {
            _tprintf(_T("CryptCreateHash failed with %X\n"), GetLastError());
            __leave;
        }

        // Hash the supplied secret password
        fResult = CryptHashData(hHash, (LPBYTE)argv[1], (DWORD)_tcslen(argv[1]), 0);

        if (!fResult) {
            _tprintf(_T("CryptHashData failed with %X\n"), GetLastError());
            __leave;
        }

        // Derive a symmetric session key from password hash
        fResult = CryptDeriveKey(hProv, CALG_RC4, hHash, 0, &hSessionKey);

        if (!fResult) {
            _tprintf(_T("CryptDeriveKey failed with %X\n"), GetLastError());
            __leave;
        }

        do {
            dwByteCount = 0;
            // Now read data from the input file 64K bytes at a time.
            fResult = ReadFile(hInFile, pbBuffer, IN_BUFFER_SIZE, &dwByteCount, NULL);

            // If the file size is exact multiple of 64K, dwByteCount will be zero after
            // all the data has been read from the input file. In this case, simply break
            // from the while loop. The check to do this is below
            if (dwByteCount == 0)
                break;

            if (!fResult) {
                _tprintf(_T("ReadFile failed with %d\n"), GetLastError());
                __leave;
            }

            finished = (dwByteCount < IN_BUFFER_SIZE);

            // Encrypt/Decrypt depending on the required action.
            if (fEncrypt) {
                fResult = CryptEncrypt(hSessionKey, 0, finished, 0, pbBuffer, &dwByteCount,
                                       OUT_BUFFER_SIZE);

                if (!fResult) {
                    _tprintf(_T("CryptEncrypt failed with %X\n"), GetLastError());
                    __leave;
                }
            } else {
                fResult = CryptDecrypt(hSessionKey, 0, finished, 0, pbBuffer, &dwByteCount);

                if (!fResult) {
                    _tprintf(_T("CryptDecrypt failed with %X\n"), GetLastError());
                    __leave;
                }
            }

            // Write the encrypted/decrypted data to the output file.
            fResult = WriteFile(hOutFile, pbBuffer, dwByteCount,
                                &dwBytesWritten, NULL);

            if (!fResult) {
                _tprintf(_T("WriteFile failed with %d\n"), GetLastError());
                __leave;
            }
        } while (!finished);

        if (fEncrypt)
            _tprintf(_T("File %s is encrypted successfully!\n"), argv[3]);
        else
            _tprintf(_T("File %s is decrypted successfully!\n"), argv[3]);
    } __finally {
        /* Cleanup */
        if (hInFile != INVALID_HANDLE_VALUE) CloseHandle(hInFile);

        if (hOutFile != INVALID_HANDLE_VALUE) CloseHandle(hOutFile);

        if (hSessionKey != NULL) CryptDestroyKey(hSessionKey);

        if (hHash != NULL) CryptDestroyHash(hHash);

        if (hProv != NULL) CryptReleaseContext(hProv, 0);
    }
}
